import numpy as np
import platgo as pg


class DTLZ2(pg.Problem):
    
    def __init__(self, M: int = 3) -> None:
        self.name = 'DTLZ2'
        self.type = '111'
        self.M = M
        self.D = M + 9
        lb = [0] * self.D
        ub = [1] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()
    
    def cal_obj(self, pop: pg.Population) -> None:
        decs = pop.decs
        XM = decs[:, (self.M - 1):]
        g = 100 * (self.D - self.M + 1 + np.sum(((XM - 0.5) ** 2 - np.cos(20 * np.pi * (XM - 0.5))), 1, keepdims=True))
        ones_matrix = np.ones((pop.N, 1))
        f = np.fliplr(np.cumprod(np.hstack([ones_matrix, np.cos(decs[:, :self.M-1] * np.pi / 2)]), 1)) * \
            np.hstack([ones_matrix, np.sin(decs[:, range(self.M - 2, -1, -1)] * np.pi / 2)]) * np.tile(1 + g, (1, self.M))
        pop.objv = f  # 把求得的目标函数值赋值给种群pop的ObjV
    
    def get_optimal(self) -> np.ndarray:
        # 目标空间均匀采样函数还未完成
        raise NotImplementedError("get optimal has not been implemented")


